package org.dreamwork.sso.saml2.protocol;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dreamwork.misc.Base64;
import org.dreamwork.secure.AlgorithmMapping;
import org.dreamwork.misc.XMLUtil;
import org.dreamwork.util.StringUtil;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.joda.time.format.ISODateTimeFormat;
import org.opensaml.xml.XMLObject;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.URLEncoder;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.util.UUID;
import java.util.zip.Deflater;
import java.util.zip.Inflater;

/**
 * Created with IntelliJ IDEA.
 * User: seth.yang
 * Date: 13-8-12
 * Time: 下午6:11
 */
public final class SamlLoginRequest extends SamlObject<XMLObject> {
    private static final String ISSUER_NODE_NAME = "Issuer";
    private static final String ENC = "utf-8";

    private String assertionConsumerServiceURL;
    private String destination;
    private String id;
    private DateTime issueInstant;
    private String protocol;
    private String issuer;

    private Log logger = LogFactory.getLog (SamlLoginRequest.class);
    private Element root;

    public SamlLoginRequest () {
        super (null);
    }

    public SamlLoginRequest (String content) {
        super (null);
        try {
            Document doc = XMLUtil.parse (content);
            root = doc.getDocumentElement ();
        } catch (Exception ex) {
            throw new RuntimeException (ex);
        }

        parse (root);
    }

    public String getAssertionConsumerServiceURL () {
        return assertionConsumerServiceURL;
    }

    public void setAssertionConsumerServiceURL (String assertionConsumerServiceURL) {
        this.assertionConsumerServiceURL = assertionConsumerServiceURL;
    }

    public String getDestination () {
        return destination;
    }

    public void setDestination (String destination) {
        this.destination = destination;
    }

    public String getId () {
        return id;
    }

    public void setId (String id) {
        this.id = id;
    }

    public DateTime getIssueInstant () {
        return issueInstant;
    }

    public void setIssueInstant (DateTime issueInstant) {
        this.issueInstant = issueInstant;
    }

    public String getProtocol () {
        return protocol;
    }

    public void setProtocol (String protocol) {
        this.protocol = protocol;
    }

    public String getIssuer () {
        return issuer;
    }

    public void setIssuer (String issuer) {
        this.issuer = issuer;
    }

    public Element getDom () {
        if (root == null) try {
            if (StringUtil.isEmpty (issuer))
                throw new RuntimeException ("issuer can't be null");

            Document doc = XMLUtil.createDocument ();
            root = doc.createElementNS (Namespace.SAML2P.uri, Namespace.SAML2P.prefix + ":AuthnRequest");
            doc.appendChild (root);

            if (StringUtil.isEmpty (assertionConsumerServiceURL))
                assertionConsumerServiceURL = "";
            root.setAttribute ("AssertionConsumerServiceURL", assertionConsumerServiceURL);
            if (StringUtil.isEmpty (destination))
                destination = "";
            root.setAttribute ("Destination", destination);
            if (StringUtil.isEmpty (id))
                id = "_" + UUID.randomUUID ();
            root.setAttribute ("ID", id);
            if (issueInstant == null)
                issueInstant = new DateTime ();
            String instant = ISODateTimeFormat.dateTime ().print (
                    issueInstant.toDateTime (DateTimeZone.UTC)
            );
            root.setAttribute ("IssueInstant", instant);
            root.setAttribute ("ProtocolBinding", "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST");
            root.setAttribute ("Version", "2.0");

            Element element = doc.createElementNS (
                    Namespace.ASSERTION.uri,
                    Namespace.ASSERTION.prefix + ":" + ISSUER_NODE_NAME
            );
            element.appendChild (doc.createTextNode (issuer));
            root.appendChild (element);
        } catch (Exception ex) {
            logger.warn (ex.getMessage (), ex);
            throw new RuntimeException (ex);
        }
        return root;
    }

    public String sign (PrivateKey privateKey, String signAlgorithmURI) {
        try {
            String temp = createRawText (signAlgorithmURI);

            // sign message
            String algorithm = AlgorithmMapping.translateSignatureAlgorithm (signAlgorithmURI);
            Signature sig = Signature.getInstance (algorithm);
            sig.initSign (privateKey);
            sig.update (temp.getBytes (ENC));
            byte[] buff = sig.sign ();
            // base64 encode signed message
            String signature = new String (Base64.encode (buff), ENC);
            if (logger.isDebugEnabled ()) {
                logger.debug ("--------- Signature-----------");
                logger.debug (signature);
                logger.debug ("------------------------------");
            }
            return signature;
        } catch (Exception ex) {
            logger.warn (ex.getMessage (), ex);
            throw new RuntimeException (ex);
        }
    }

    public boolean validate (PublicKey publicKey, String encodedSignature, String algorithmURI) {
        try {
            // verify saml request
            String temp = createRawText (algorithmURI);
            byte[] buff = Base64.decode (encodedSignature.replace(" ", "+"));

            String algorithmName = AlgorithmMapping.translateSignatureAlgorithm (algorithmURI);
            Signature signature = Signature.getInstance (algorithmName);
            signature.initVerify (publicKey);
            signature.update (temp.getBytes (ENC));
            return signature.verify (buff);
        } catch (Exception ex) {
            logger.warn (ex.getMessage (), ex);
            throw new RuntimeException (ex);
        }
    }

    private void parse (Element root) {
        assertionConsumerServiceURL = root.getAttribute ("AssertionConsumerServiceURL");
        destination = root.getAttribute ("Destination");
        id = root.getAttribute ("ID");
        String s_issueInstant = root.getAttribute ("IssueInstant");
        issueInstant = ISODateTimeFormat.dateTime ().parseDateTime (s_issueInstant);
        protocol = root.getAttribute ("ProtocolBinding");

        NodeList list = root.getChildNodes ();
        if (list == null)
            throw new RuntimeException ("Can't find issuer!!!");

        for (int i = 0; i < list.getLength (); i ++) {
            Node node = list.item (i);
            if (node.getNodeType () != Node.ELEMENT_NODE) continue;

            String localName = node.getLocalName ();
            String namespaceURI = node.getNamespaceURI ();

            if (localName == null || !localName.equals (ISSUER_NODE_NAME))
                continue;

            if (namespaceURI == null || !namespaceURI.equals (Namespace.ASSERTION.uri))
                throw new RuntimeException ("Unknown namespace: " + namespaceURI);

            issuer = node.getTextContent ();
        }
    }

    public String createRawText (String signAlgorithm) throws IOException {
        // deflate raw SAML2 Request
        byte[] rawData = getXMLBytes ();
        Deflater deflater = new Deflater (8, true);
        deflater.setInput (rawData);
        byte[] buff = new byte[rawData.length];
        deflater.finish ();
        int length = deflater.deflate (buff);
        byte[] a = new byte[length];
        System.arraycopy (buff, 0, a, 0, length);

        // base64 encode deflated data
        String samlRequest = new String (Base64.encode (a), ENC);
        // url encode base64 encoded message
        samlRequest = URLEncoder.encode (samlRequest, ENC);

        // signature algorithm uri
        String sigAlg = URLEncoder.encode (signAlgorithm, ENC);

        // prepare the message for signing
        String temp = "SAMLRequest=" + samlRequest + "&SigAlg=" + sigAlg;
        if (logger.isDebugEnabled ()) {
            logger.debug ("----------------- row string -----------------");
            logger.debug (temp);
            logger.debug ("----------------------------------------------");
        }
        return temp;
    }

    public static SamlLoginRequest parse (String encryptedText) {
        try {
            // decode saml request
            byte[] buff = Base64.decode (encryptedText.getBytes (ENC));
            Inflater inflater = new Inflater (true);
            inflater.setInput (buff);
            byte[] result = new byte[1024];
            ByteArrayOutputStream baos = new ByteArrayOutputStream ();
            while (!inflater.finished ()) {
                int length = inflater.inflate (result);
                baos.write (result, 0, length);
            }

            // build saml request
            return new SamlLoginRequest (baos.toString (ENC));
        } catch (Exception ex) {
            LogFactory.getLog (SamlLoginRequest.class).warn (ex.getMessage (), ex);
            throw new RuntimeException (ex);
        }
    }
}